#! /usr/bin/env python
# coding=utf-8

# 在目标物种中获得位数 需要输入 趋同位点的位点号 
# 对应的 cut_loci 文件 
# 原始的比对文件

# 同时截取 修改后的文件 
# 获得对应物种 该位置的 pep 和 mRNA 比对 上下游 10 bp
import sys

fra = 10
import sys
import argparse




dis = '''
切取 目标物种的 趋同位点上下游各10位序列 这个序列位于中间
前后氨基酸 没有达到 10位的不会被选取
tblastn_fetch_gene01_take_ref_target_site.py -r pep_raw_ali.fa -m filter_add_mRNA_ali.fa -p filter_add_pep_ali.fa -f cut -s Nig,Hte -c cut_loci -t 6
'''


parser = argparse.ArgumentParser(description=dis)



parser.add_argument('-r',
                help='必须给定，参考物种的原始比对')

parser.add_argument('-p',
                help='必须给定，当前比对的蛋白序列')

parser.add_argument('-m',
                help='必须给定，当前比对的mRNA')

parser.add_argument('-f',
                help='必须给定，前缀')

parser.add_argument('-s',
                help='必须给定，参考物种')


parser.add_argument('-t',
                help='必须给定，tblastn 结果路径')

parser.add_argument('-c',
                help='必须给定，目标物种列表文件')




# 获取参数
args = parser.parse_args()



if not args.p or not args.r or not args.m or not args.t or not args.f or not args.c or not args.s:
    parser.print_help()
    sys.exit()




infile_raw_pep = args.r  # r'pep_raw_ali.fa'

infile_ali_pep = args.p #'filter_add_pep_ali.fa'

infile_ali_mRNA = args.m  #r'filter_add_mRNA_ali.fa'

 
prefix = args.f # 'cut'

outfile_ali_mRNA = open(prefix+'_mRNA_ali.fa','w')

outfile_ali_pep = open(prefix+ '_pep_ali.fa','w')



cut_loci = args.c

ref_spe = args.s.split(',')

target_site = int(args.t)


target_loci = None

from Bio import SeqIO




with open(cut_loci) as fila:
	for i in fila:
		k= i.strip().split('\t')
		if k[0] == str(target_site+1):
			target_loci = int(k[2])-1

if target_loci == None:
	print('loci error')
	sys.exit(666)


# 在原始的pep 序列中 提取目标位点的下游序列
query_seq = None
query_seq_dic = {}
for i in SeqIO.parse(infile_raw_pep,'fasta'):
	# 这里必须要求 原始的 pep文件 其id 就是物种缩写名


	if str(i.name).split(' ')[0] in ref_spe:
		#print(i)
		query_seq = str(i.seq)[target_loci:target_loci+30].replace('-','')
		if query_seq == None or len(query_seq) <8:
			pass
		else:
			query_seq_dic[str(i.name)] = query_seq

if query_seq_dic == {}:
	print('not enough query len')
	sys.exit(666)

# 使用该下游序列为query 搜索目标


# 找到新命名的第三项
import re
r = re.compile('[A-Z][a-z]+[A-Z][a-z]+([A-Z][a-z]+)')



current_loci = None
raw_pep_dic = {}
raw_mRNA_dic = {}
for i in SeqIO.parse(infile_ali_pep,'fasta'):
	t = r.findall(i.name)
	if t!=[] and t[0] in query_seq_dic:
		l = 0
		s = str(i.seq).replace('-','').find(query_seq_dic[t[0]])

		if s != -1:
			c  = 0
			while c < s:
				if i.seq[l] != '-':
					c += 1
				l+= 1 
		current_loci = l
		
	raw_pep_dic[str(i.name)] = str(i.seq)

for i in SeqIO.parse(infile_ali_mRNA,'fasta'):
	raw_mRNA_dic[str(i.name)] = str(i.seq)


cut_pep_dic = {}
if current_loci!= None:
	for i in raw_pep_dic:
		seq_before = raw_pep_dic[i][max([current_loci-fra,0]):current_loci]
		mRNA_before = raw_mRNA_dic[i][max([current_loci*3-fra*3,0]):current_loci*3]
		if current_loci-fra < 0:
			seq_before = '-'*(-(current_loci-fra)) + seq_before
			mRNA_before = '---'*(-(current_loci-fra)) + mRNA_before
		seq_after = raw_pep_dic[i][current_loci:current_loci+11]
		mRNA_after = raw_mRNA_dic[i][current_loci*3 : current_loci*3 + 33]
		if len(seq_after) <= 10:
			seq_after += '-'*(10 - len(seq_after))
			mRNA_after = '---'*(10 - len(seq_after))

		seq = seq_before+seq_after
		mRNA = mRNA_before+mRNA_after	


		#cut_seq = rcurrent_loci:10]

		if mRNA.replace('-','') == '':
			continue

		outfile_ali_mRNA.write('>'+i+'\n'+mRNA+'\n')
		outfile_ali_pep.write('>'+i+'\n'+seq+'\n')


outfile_ali_mRNA.close()
outfile_ali_pep.close()


